import os
import random
from shutil import copyfile

# 数据集根目录
data_dir = "data/mydata"

# 创建分割后的训练集和验证集文件夹
train_img_dir = os.path.join(data_dir, "train_images")
train_label_dir = os.path.join(data_dir, "train_labels")
val_img_dir = os.path.join(data_dir, "val_images")
val_label_dir = os.path.join(data_dir, "val_labels")
os.makedirs(train_img_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_img_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)

# 获取所有图像文件和标签文件的路径
image_dir = os.path.join(data_dir, "images")
label_dir = os.path.join(data_dir, "labels")
image_files = [img for img in os.listdir(image_dir) if img.endswith(".jpg")]
random.shuffle(image_files)  # 随机打乱图像文件顺序

# 分割数据集为训练集和验证集
split_ratio = 0.8  # 训练集和验证集的比例
split_index = int(len(image_files) * split_ratio)
train_image_files = image_files[:split_index]
val_image_files = image_files[split_index:]

# 将训练图像文件复制到训练集图片文件夹，并生成对应的标签文件
for img in train_image_files:
    src_img = os.path.join(image_dir, img)
    dst_img = os.path.join(train_img_dir, img)
    copyfile(src_img, dst_img)

    label_file = os.path.splitext(img)[0] + ".txt"
    src_label = os.path.join(label_dir, label_file)
    dst_label = os.path.join(train_label_dir, label_file)
    copyfile(src_label, dst_label)

# 将验证图像文件复制到验证集图片文件夹，并生成对应的标签文件
for img in val_image_files:
    src_img = os.path.join(image_dir, img)
    dst_img = os.path.join(val_img_dir, img)
    copyfile(src_img, dst_img)

    label_file = os.path.splitext(img)[0] + ".txt"
    src_label = os.path.join(label_dir, label_file)
    dst_label = os.path.join(val_label_dir, label_file)
    copyfile(src_label, dst_label)

# 生成 train.txt 文件
train_txt_path = os.path.join(data_dir, "train.txt")
with open(train_txt_path, "w") as f:
    for img in train_image_files:
        img_path = os.path.join("train_images", img)
        label_path = os.path.join("train_labels", os.path.splitext(img)[0] + ".txt")
        f.write(f"{img_path} {label_path}\n")


